from patsy import dmatrices, dmatrix, demo_data
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import statsmodels.api as sm
from sklearn.metrics import mean_squared_error, mean_absolute_error


class SplineRegression(object):
    """
    参数
    ----------
    train_x ：bool 训练数据, ndarray数组
    spline_type : str 样条类型：B-样条(bs)，自然样条(cr)，默认为bs
    df ：int 自由度 df = knots + 4-1，默认为7
    knots : tuple 节点，例如(10,21,33)
    degree : int 次方数，默认为3，即三次样条
    include_intercept : bool 截距，默认是不需要截距
    return_type : str 返回类型 默认为dataframe
    """
    def __init__(self, _train_x, spline_type="bs", df=7, knots=(0, 10), degree=3,
                 include_intercept=False, return_type='dataframe'):

        self.train_x = _train_x
        self.spline_type = spline_type
        self.df = df
        self.knots = knots
        self.degree = degree
        self.include_intercept = include_intercept
        self.return_type = return_type
        self.transformed_x = self.set_dmatrix(self.train_x)
        self._y_predict = None

    def set_dmatrix(self, _train_x):
        """
        设置样条函数
        return 返回设置好的样条函数
        """
        _first_para = str(self.spline_type) + "(train,df=" + str(self.df) + ",knots=" + \
                     str(self.knots) + ",degree=" + str(self.degree) + ", include_intercept=" +\
                     str(self.include_intercept) + ")"
        _second_para = {"train": _train_x}
        return dmatrix(_first_para, _second_para, return_type=self.return_type)

    def fit(self, y_train):
        self.fitted = sm.GLM(y_train, self.transformed_x).fit()

    def predict(self, x_test):
        _data = self.set_dmatrix(x_test)
        self._y_predict = self.fitted.predict(_data)
        return self._y_predict

    def get_mse(self, y_test):
        """计算:MSE"""
        return mean_squared_error(y_test, self._y_predict)

    def get_rmse(self, y_test):
        """计算:RMSE"""
        return np.sqrt(mean_squared_error(y_test, self._y_predict))

    def get_mae(self, y_test):
        """计算:MAE"""
        return mean_absolute_error(y_test, self._y_predict)


if __name__ == '__main__':
    file_path = "../test_file/data.csv"
    data = pd.read_csv(file_path)
    data_x, data_y = data['times'], data['accel']
    train_x, valid_x, train_y, valid_y = train_test_split(data_x, data_y, test_size=0.33, random_state=1)

    spline_type = "bs"
    knots1 = (10, 32, 40)
    df = 6
    degree = 4
    knots2 = (15, 18, 32, 40)

    sr1 = SplineRegression(train_x, df=df, knots=knots1)
    sr1.fit(train_y)
    pred1 = sr1.predict(valid_x)

    sr2 = SplineRegression(train_x, knots=knots2)
    sr2.fit(train_y)
    pred2 = sr2.predict(valid_x)

    rmse1 = sr1.get_rmse(valid_y)
    rmse2 = sr2.get_rmse(valid_y)

    # 用70个观察值画图
    xp = np.linspace(valid_x.min(), valid_x.max(), 70)
    # 预测
    pred11 = sr1.predict(xp)
    pred12 = sr2.predict(xp)

    # 绘制样条曲线和误差曲线
    plt.scatter(data.times, data.accel, facecolor='None', edgecolor='k', alpha=0.1)
    plt.plot(xp, pred11, label='Specifying degree =3 with 3 knots')
    plt.plot(xp, pred12, color='r', label='Specifying degree =3 with 4 knots')
    plt.legend()
    plt.xlabel('times')
    plt.ylabel('accel')
    plt.show()





